#将split_train_val_yolo.py生成的train.txt val.txt两个文件进行读取，
#并查询对应json文件的backgroundNumber信息,存储成txt文件
import argparse
from pathlib import Path
import xml.etree.ElementTree as ET

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="generateBackgroundTxt")
    parser.add_argument('--root-dir', default="", type=str)
    return parser.parse_args()

def main():
    args = Parse_Arguments()
    root_dir = args.root_dir
    file_path = []
    with open(root_dir + '/dataSet_path/train.txt', 'r') as f:
        file_path += f.read().strip().splitlines()
    with open(root_dir + '/dataSet_path/val.txt', 'r') as f:
        file_path += f.read().strip().splitlines()

    txt_file_content = []
    for item in file_path:
        xml_filename = Path(item).with_suffix('.xml').name
        xml_file_path = root_dir + f"/Annotations/{xml_filename}"
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        try:
            value = root.find('backgroundNumber').text
            txt_file_content.append(f"{item} {value}")
        except:
            print("发生错误！！！")
            print(f"有可能是该{xml_filename}里没有加入backgroundNumber标签")
    with open(root_dir + '/dataSet_path/residual.txt', 'w') as f:
        for line in txt_file_content:
            f.write(line + '\n')
if __name__ == "__main__":
    main()
